import datetime

import gymnasium as gym
import gymnasium.wrappers as gym_wrap
import matplotlib.pyplot as plt
import numpy as np
import torch

import DQN_model as DQN

env = gym.make("CarRacing-v2", continuous=False) # render_mode="human"
env = DQN.SkipFrame(env, skip=4)
env = gym_wrap.GrayScaleObservation(env)
env = gym_wrap.ResizeObservation(env, shape=84)
env = gym_wrap.FrameStack(env, num_stack=4)
state, info = env.reset()
action_n = env.action_space.n
driver = DQN.Agent(state.shape, action_n, double_q=True)
batch_n = 32
play_n_episodes = 3000
episode_epsilon_list = []
episode_reward_list = []
episode_length_list = []
episode_loss_list = []
episode_date_list = []
episode_time_list = []
episode = 0
timestep_n = 0
when2learn = 4  # in timesteps
when2sync = 5000  # in timesteps
when2save = 100000  # in timesteps
when2report = 5000  # in timesteps
when2eval = 50000  # in timesteps
when2log = 10  # in episodes
report_type = 'plot'  # 'text', 'plot', None

while episode <= play_n_episodes:
    episode += 1
    episode_reward = 0
    episode_length = 0
    updating = True
    loss_list = []
    episode_epsilon_list.append(driver.epsilon)

    while updating:
        timestep_n += 1
        episode_length += 1

        action = driver.take_action(state)
        new_state, reward, terminated, truncated, info = env.step(action)
        episode_reward += reward
        driver.store(state, action, reward, new_state, terminated)
        # Move to the next state
        state = new_state
        updating = not (terminated or truncated)

        if timestep_n % when2sync == 0:
            upd_net_param = driver.updating_net.state_dict()
            driver.frozen_net.load_state_dict(upd_net_param)

        if timestep_n % when2save == 0:
            driver.save(driver.save_dir, 'DQN')

        if timestep_n % when2learn == 0:
            q, loss = driver.update_net(batch_n)
            loss_list.append(loss)

        if timestep_n % when2report == 0 and report_type == 'text':
            print(f'Report: {timestep_n} timestep')
            print(f'    episodes: {episode}')
            print(f'    n_updates: {driver.n_updates}')
            print(f'    epsilon: {driver.epsilon}')

        if timestep_n % when2eval == 0 and report_type == 'text':
            rewards_tensor = torch.tensor(episode_reward_list,
                                          dtype=torch.float)
            eval_reward = torch.clone(rewards_tensor[-50:])
            mean_eval_reward = round(torch.mean(eval_reward).item(), 2)
            std_eval_reward = round(torch.std(eval_reward).item(), 2)

            lengths_tensor = torch.tensor(episode_length_list,
                                          dtype=torch.float)
            eval_length = torch.clone(lengths_tensor[-50:])
            mean_eval_length = round(torch.mean(eval_length).item(), 2)
            std_eval_length = round(torch.std(eval_length).item(), 2)

            print(f'Evaluation: {timestep_n} timestep')
            print(f'    reward {mean_eval_reward}±{std_eval_reward}')
            print(f'    episode length {mean_eval_length}±{std_eval_length}')
            print(f'    episodes: {episode}')
            print(f'    n_updates: {driver.n_updates}')
            print(f'    epsilon: {driver.epsilon}')

    state, info = env.reset()

    episode_reward_list.append(episode_reward)
    episode_length_list.append(episode_length)
    episode_loss_list.append(np.mean(loss_list))
    now_time = datetime.datetime.now()
    episode_date_list.append(now_time.date().strftime('%Y-%m-%d'))
    episode_time_list.append(now_time.time().strftime('%H:%M:%S'))

    if report_type == 'plot':
        draw_check = DQN.plot_reward(episode, episode_reward_list, timestep_n)

    if episode % when2log == 0:
        driver.write_log(
            episode_date_list,
            episode_time_list,
            episode_reward_list,
            episode_length_list,
            episode_loss_list,
            episode_epsilon_list,
            log_filename='DQN_log_test.csv'
        )

if report_type == 'text':
    rewards_tensor = torch.tensor(episode_reward_list, dtype=torch.float)
    eval_reward = torch.clone(rewards_tensor[-100:])
    mean_eval_reward = round(torch.mean(eval_reward).item(), 2)
    std_eval_reward = round(torch.std(eval_reward).item(), 2)

    lengths_tensor = torch.tensor(episode_length_list, dtype=torch.float)
    eval_length = torch.clone(lengths_tensor[-100:])
    mean_eval_length = round(torch.mean(eval_length).item(), 2)
    std_eval_length = round(torch.std(eval_length).item(), 2)

    print(f'Final evaluation: {timestep_n} timestep')
    print(f'    reward {mean_eval_reward}±{std_eval_reward}')
    print(f'    episode length {mean_eval_length}±{std_eval_length}')
    print(f'    episodes: {episode}')
    print(f'    n_updates: {driver.n_updates}')
    print(f'    epsilon: {driver.epsilon}')

driver.save(driver.save_dir, 'DQN')
driver.write_log(
    episode_date_list,
    episode_time_list,
    episode_reward_list,
    episode_length_list,
    episode_loss_list,
    episode_epsilon_list,
    log_filename='DQN_log_test.csv'
)
env.close()
plt.ioff()
plt.show()

